"""Data collection and management module for Quafu quantum platform.

This module handles experiment data collection, submission and result processing
for the Quafu quantum computing platform.
"""

from datetime import datetime
from typing import List, Optional

import pandas as pd
from quafu import Task
from quafu import QuantumCircuit
from .quafu_platform import (
    submit_circ_jobs, 
    get_result_less_15qubits, 
    get_result_more_15qubits,
)


class QuafuData:
    """Handler for Quafu platform experiment data.
    
    This class manages the submission, collection and storage of quantum experiment
    data from the Quafu platform.
    
    Attributes:
        num_qubits: Number of qubits in experiments
        num_shots: Number of shots per experiment
        platform: Quafu platform instance
        exp_name: Experiment name with timestamp
        df: DataFrame storing experiment results
    """

    def __init__(self, 
                 num_qubits: int,
                 num_shots: int,
                 platform: Task,
                 exp_name: str):
        """Initialize the data handler.

        Args:
            num_qubits: Number of qubits in experiments
            num_shots: Number of shots per experiment
            platform: Quafu platform instance
            exp_name: Base name for experiments
        """
        self.num_qubits = num_qubits
        self.num_shots = num_shots
        self.platform = platform
        self.exp_name = f"{exp_name}_{datetime.now().strftime('%Y%m%d%H%M%S')}"
        
        # Initialize results dataframe
        self.df = pd.DataFrame(columns=[
            'lab_id', 'exp_id', 'exp_name', 'qubits', 
            'true_value', 'exp_result', 'circuit'
        ])
    
    def submit_experiment(self, circuit_list: List[QuantumCircuit]) -> None:
        """Submit quantum circuits for execution.

        Args:
            circuit_list: List of quantum circuits to execute
        """
        submit_circ_jobs(
            circuit_list, 
            self.df, 
            self.platform, 
            self.exp_name
        )

    def receive_result(self) -> None:
        """Collect experiment results from platform."""
        if self.num_qubits < 15:
            get_result_less_15qubits(self.df, self.platform)
        else:
            get_result_more_15qubits(self.df, self.platform)
    
    def save_result(self, file_name: str) -> None:
        """Save results to CSV file.

        Args:
            file_name: Path to output CSV file
        """
        self.df.to_csv(file_name, index=False)
    